import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os
import cv2
import h5py
from models.convnext.convnext import ConvNeXt

def predictSingleImage(imagePath, weight):
    print(imagePath)
    print(weight)

    categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch',
                  'Pattern_fail',
                  'PR_peeling', 'Seam', 'Reference', 'Surface_particle', 'Burried_particle', 'Cu_diffuse',
                  'Prelayer_defect_developed',
                  'Void', 'Residue', 'Scratch']

    categories_to_id = dict((c, i) for i, c in enumerate(categories))
    id_to_categories = dict((v, k) for k, v in categories_to_id.items())


    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    model = torch.load(weight)
    model.eval()
    model.to(DEVICE)

    # preprocess
    transform_test = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.40272543], std=[0.13901867])
        ])
    f = h5py.File(imagePath, 'r')  # 打开h5文件
    class_tiff = f['class_tiff'][:]
    defect_tiff = f['defect_tiff'][:]
    f.close()

    image_class = Image.fromarray(class_tiff).convert("L")
    image_class = transform_test(image_class)
    image_class = torch.unsqueeze(image_class, dim=0).to(DEVICE)

    # Predict
    out = model(image_class)
    _, pred = torch.max(out.data, 1)
    res = id_to_categories[pred.data.item()] # defect name
    # print('Image Name:{} \npredict:{}'.format(imagepath, id_to_categories[pred.data.item()]))
    return res

def getStatedict(origin, output):
    model = torch.load(origin)
    torch.save(model.state_dict(), output)

def predictTestDIr(testDirPath, weight):
    categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch',
                  'Pattern_fail',
                  'PR_peeling', 'Seam', 'Reference', 'Surface_particle', 'Burried_particle', 'Cu_diffuse',
                  'Prelayer_defect_developed',
                  'Void', 'Residue', 'Scratch']

    categories_to_id = dict((c, i) for i, c in enumerate(categories))
    id_to_categories = dict((v, k) for k, v in categories_to_id.items())

    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = torch.load(weight)
    model.eval()
    model.to(DEVICE)

    # preprocess
    transform_test = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.40272543], std=[0.13901867])
        ])

    count = 0
    right = 0

    for categoriy in os.listdir(testDirPath):
        file_lists = os.listdir(os.path.join(testDirPath, categoriy))
        categoriy_count = 0
        categoriy_right = 0
        for file in file_lists:
            count += 1
            categoriy_count += 1
            imagePath = os.path.join(testDirPath, categoriy, file)
            image = None
            f = h5py.File(imagePath, 'r')  # 打开h5文件
            class_tiff = f['class_tiff'][:]
            defect_tiff = f['defect_tiff'][:]
            f.close()
            image_class = Image.fromarray(class_tiff).convert("L")
            image_class = transform_test(image_class)
            image_class = torch.unsqueeze(image_class, dim=0).to(DEVICE)

            # Predict
            out = model(image_class)
            _, pred = torch.max(out.data, 1)
            res = id_to_categories[pred.data.item()] # defect name

            if(res == categoriy):
                right += 1
                categoriy_right += 1
            else:
                a = None
                # print('Image Name:{} \npredict:{}\ncategory:{}'.format(file, id_to_categories[pred.data.item()], categoriy))
        print("{}:{} \n acc:{}".format(categoriy, categoriy_count, categoriy_right / categoriy_count))
        print("======================================================================")
    print("total:{} \n acc:{}".format(count, right / count))




if __name__ == '__main__':
    imagepath = "/home/yeadc/Documents/cxj/datasets_h5_v2/val/Burried_particle/Bump_FY-C201F_MD_ASI_lFFK100951_w1_27680_Topography3.h5"
    weight = '/home/yeadc/Documents/cxj/Classify-Seg/task_classify/task/model_167_90.672.pth'


    # testDir = "/home/yeadc/Documents/cxj/datasets_h5_v2/val"
    #
    # predictTestDIr(testDir, weight)
    predictSingleImage(imagepath, weight)

'''
swinv2:  acc:0.9366678661388989
convnext_tiny:  acc:0.9093198992443325
'''
